"""
MACD策略

基于MACD指标的趋势跟踪策略。

作者: AI Assistant
版本: 1.0.0
日期: 2025-01-06
"""

import pandas as pd
import numpy as np
from backtest.strategy import StrategyBase


class MACDStrategy(StrategyBase):
    """
    MACD策略
    
    使用MACD柱状图的正负变化作为买卖信号：
    - MACD柱由负转正（金叉）：买入
    - MACD柱由正转负（死叉）：卖出
    
    参数:
        fast_period: 快线周期，默认12
        slow_period: 慢线周期，默认26
        signal_period: 信号线周期，默认9
        
    信号规则:
        - MACD柱由负转正：买入
        - MACD柱由正转负：卖出
        - 其他情况：持有
        
    示例:
        >>> strategy = MACDStrategy(fast_period=12, slow_period=26, signal_period=9)
        >>> engine = BacktestEngine(strategy=strategy, data=data)
        >>> result = engine.run()
    """
    
    def __init__(
        self,
        fast_period: int = 12,
        slow_period: int = 26,
        signal_period: int = 9
    ):
        """
        初始化MACD策略
        
        参数:
            fast_period: 快线EMA周期
            slow_period: 慢线EMA周期
            signal_period: 信号线EMA周期
        """
        if fast_period >= slow_period:
            raise ValueError("快线周期必须小于慢线周期")
        
        if fast_period < 2 or slow_period < 2 or signal_period < 2:
            raise ValueError("所有周期必须大于1")
        
        super().__init__(
            name=f"MACD策略({fast_period},{slow_period},{signal_period})",
            min_period=slow_period + signal_period,
            fast_period=fast_period,
            slow_period=slow_period,
            signal_period=signal_period
        )
        
        self.fast_period = fast_period
        self.slow_period = slow_period
        self.signal_period = signal_period
        self.last_signal = None
    
    def _calculate_macd(self, data: pd.DataFrame) -> tuple:
        """
        计算MACD指标
        
        参数:
            data: 历史数据
            
        返回:
            tuple: (DIF, DEA, MACD柱)
        """
        close = data['收盘']
        
        # 计算快线和慢线EMA
        ema_fast = close.ewm(span=self.fast_period, adjust=False).mean()
        ema_slow = close.ewm(span=self.slow_period, adjust=False).mean()
        
        # 计算DIF（快线-慢线）
        dif = ema_fast - ema_slow
        
        # 计算DEA（DIF的EMA）
        dea = dif.ewm(span=self.signal_period, adjust=False).mean()
        
        # 计算MACD柱（DIF-DEA）* 2
        macd = (dif - dea) * 2
        
        return dif, dea, macd
    
    def generate_signal(self, data: pd.DataFrame, current_bar: pd.Series) -> str:
        """
        生成交易信号
        
        参数:
            data: 当前及之前的所有数据
            current_bar: 当前K线数据
            
        返回:
            str: 'BUY', 'SELL', 或 'HOLD'
        """
        # 确保有足够的数据
        if len(data) < self.min_period:
            return 'HOLD'
        
        # 计算MACD
        dif, dea, macd = self._calculate_macd(data)
        
        # 需要至少两个MACD值来判断交叉
        if len(macd) < 2:
            return 'HOLD'
        
        # 获取当前和前一天的MACD柱值
        macd_current = macd.iloc[-1]
        macd_prev = macd.iloc[-2]
        
        # 检查是否有NaN值
        if pd.isna(macd_current) or pd.isna(macd_prev):
            return 'HOLD'
        
        # 金叉：MACD柱由负转正
        if macd_prev < 0 and macd_current > 0:
            if self.last_signal != 'BUY':
                self.last_signal = 'BUY'
                return 'BUY'
        
        # 死叉：MACD柱由正转负
        elif macd_prev > 0 and macd_current < 0:
            if self.last_signal != 'SELL':
                self.last_signal = 'SELL'
                return 'SELL'
        
        return 'HOLD'
    
    def get_indicator_values(self, data: pd.DataFrame) -> pd.DataFrame:
        """
        获取策略指标值（用于绘图）
        
        参数:
            data: 历史数据
            
        返回:
            pd.DataFrame: 包含MACD指标的数据框
        """
        result = data.copy()
        dif, dea, macd = self._calculate_macd(data)
        result['DIF'] = dif
        result['DEA'] = dea
        result['MACD'] = macd
        return result

